import torch
from torch import nn
import torch.nn.init as init


class Gru(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers):
        super(Gru, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        # self.fc1 = nn.Sequential(nn.Linear(input_size, hidden_size), nn.ReLU(), nn.Linear(hidden_size, hidden_size))
        self.gru = nn.GRU(input_size, hidden_size, num_layers, batch_first=True)
        self.fc2 = nn.Sequential(nn.Linear(hidden_size, hidden_size), nn.ReLU(), nn.Linear(hidden_size, output_size))
        self.output_size = output_size


    def forward(self, x):
        # x = self.fc1(x)
        # 通过GRU层传递输入和隐藏状态
        out, _ = self.gru(x)
        # 只使用序列的最后一个输出
        out = self.fc2(out[:, -1, :])
        # 如果是单变量预测，直接就返回值
        if self.output_size == 1:
            out = out[:, 0]
        return out
    

    def predict(self, x):
        out, _ = self.gru(x)
        print("gru:")
        print(out.shape)
        print(out)
        print(_.shape)
        print(_)
        print("gru wieght")
        print(self.gru)
        print("gru end")
        out = self.fc(out[-1, :])
        return out
    

    def create_dataset(self, x):
        x = self.fc1(x)
        out, hidden = self.gru(x)
        return hidden[-1]
